
import torch
import skimage.io


from models.decoder import Decoder
from models.encoder import Encoder

class CaptionModel(object):
    def __init__(self, decoder_file, resnet101_file="./resnet101.pth", device="cpu"):
        self.encoder = Encoder(resnet101_file)
        print("====> loading checkpoint '{}'".format(decoder_file))
        chkpoint = torch.load(decoder_file, map_location=lambda s, l: s)
        decoder = Decoder(chkpoint['idx2word'], chkpoint['settings'])
        decoder.load_state_dict(chkpoint['model'])
        decoder.to(device)
        decoder.eval()
        self.device = device
        self.decoder = decoder
        
    def predict(self, image, **kwargs):
        if isinstance(image, str):
            img = skimage.io.imread(image)
        else:
            img = image
        with torch.no_grad():
            img = self.encoder.preprocess(img)
            img = img.to(self.device)
            img_feat, _ = self.encoder(img)
        rest, _ = self.decoder.sample(img_feat, **kwargs)
        return rest

if __name__ == "__main__":
    model_rl_grid = CaptionModel(decoder_file="rl/grid-based/model_49_5.1001_0418-1216.pth")
    model_rl_region = CaptionModel(decoder_file="rl/region-based/model_49_5.0719_0420-1839.pth")
    model_xe_grid = CaptionModel(decoder_file="xe/grid-based/model_49_2.8938_0417-1245.pth")
    model_xe_region = CaptionModel(decoder_file="xe/region-based/model_49_2.8568_0420-1657.pth")

    print(model_rl_grid.predict("COCO_val2014_000000391895.jpg"))
    print(model_rl_region.predict("COCO_val2014_000000391895.jpg"))
    print(model_xe_grid.predict("COCO_val2014_000000391895.jpg"))
    print(model_xe_region.predict("COCO_val2014_000000391895.jpg"))